package org.finesys.chat.core.base.json;

import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import org.finesys.ai.core.chat.json.JsonArraySchema;
import org.finesys.ai.core.chat.json.JsonBooleanSchema;
import org.finesys.ai.core.chat.json.JsonEnumSchema;
import org.finesys.ai.core.chat.json.JsonIntegerSchema;
import org.finesys.ai.core.chat.json.JsonNumberSchema;
import org.finesys.ai.core.chat.json.JsonObjectSchema;
import org.finesys.ai.core.chat.json.JsonReferenceSchema;
import org.finesys.ai.core.chat.json.JsonSchemaElement;
import org.finesys.ai.core.chat.json.JsonStringSchema;
import org.finesys.chat.core.base.agent.Description;
import org.finesys.common.core.util.Utils;

public class JsonSchemaElementHelper {

    public JsonSchemaElementHelper() {
    }

    public static JsonSchemaElement jsonSchemaElementFrom(Class<?> clazz, Type type, String fieldDescription, Map<Class<?>, VisitedClassMetadata> visited) {
        if (TypeUtils.isJsonString(clazz)) {
            return JsonStringSchema.builder().description(fieldDescription).build();
        } else if (TypeUtils.isJsonInteger(clazz)) {
            return JsonIntegerSchema.builder().description(fieldDescription).build();
        } else if (TypeUtils.isJsonNumber(clazz)) {
            return JsonNumberSchema.builder().description(fieldDescription).build();
        } else if (TypeUtils.isJsonBoolean(clazz)) {
            return JsonBooleanSchema.builder().description(fieldDescription).build();
        } else if (clazz.isEnum()) {
            return JsonEnumSchema.builder().enumValues(Arrays.stream(clazz.getEnumConstants()).map(Object::toString).collect(Collectors.toList())).description((String) Optional.ofNullable(fieldDescription).orElse(descriptionFrom(clazz))).build();
        } else if (clazz.isArray()) {
            return JsonArraySchema.builder().items(jsonSchemaElementFrom(clazz.getComponentType(), (Type) null, (String) null, visited)).description(fieldDescription).build();
        } else {
            return (JsonSchemaElement) (Collection.class.isAssignableFrom(clazz) ? JsonArraySchema.builder().items(jsonSchemaElementFrom(getActualType(type), (Type) null, (String) null, visited)).description(fieldDescription).build() : jsonObjectOrReferenceSchemaFrom(clazz, fieldDescription, visited, false));
        }
    }

    public static JsonSchemaElement jsonObjectOrReferenceSchemaFrom(Class<?> type, String description, Map<Class<?>, VisitedClassMetadata> visited, boolean setDefinitions) {
        if (visited.containsKey(type) && isCustomClass(type)) {
            VisitedClassMetadata visitedClassMetadata = (VisitedClassMetadata) visited.get(type);
            JsonSchemaElement jsonSchemaElement = visitedClassMetadata.jsonSchemaElement;
            if (jsonSchemaElement instanceof JsonReferenceSchema) {
                visitedClassMetadata.recursionDetected = true;
            }

            return jsonSchemaElement;
        } else {
            String reference = Utils.generateUUIDFrom(type.getName());
            JsonReferenceSchema jsonReferenceSchema = JsonReferenceSchema.builder().reference(reference).build();
            visited.put(type, new VisitedClassMetadata(jsonReferenceSchema, reference, false));
            Map<String, JsonSchemaElement> properties = new LinkedHashMap<>();
            Field[] var7 = type.getDeclaredFields();
            int var8 = var7.length;

            for (int var9 = 0; var9 < var8; ++var9) {
                Field field = var7[var9];
                String fieldName = field.getName();
                if (!Modifier.isStatic(field.getModifiers()) && !fieldName.equals("__$hits$__") && !fieldName.startsWith("this$")) {
                    String fieldDescription = descriptionFrom(field);
                    JsonSchemaElement jsonSchemaElement = jsonSchemaElementFrom(field.getType(), field.getGenericType(), fieldDescription, visited);
                    properties.put(fieldName, jsonSchemaElement);
                }
            }

            JsonObjectSchema.Builder builder = JsonObjectSchema.builder().description((String) Optional.ofNullable(description).orElse(descriptionFrom(type))).properties(properties).required(new ArrayList<>(properties.keySet()));
            ((VisitedClassMetadata) visited.get(type)).jsonSchemaElement = builder.build();
            if (setDefinitions) {
                Map<String, JsonSchemaElement> definitions = new LinkedHashMap<>();
                visited.forEach((clazz, visitedClassMetadatax) -> {
                    if (visitedClassMetadatax.recursionDetected) {
                        definitions.put(visitedClassMetadatax.reference, visitedClassMetadatax.jsonSchemaElement);
                    }

                });
                if (!definitions.isEmpty()) {
                    builder.definitions(definitions);
                }
            }

            return builder.build();
        }
    }

    private static String descriptionFrom(Field field) {
        return descriptionFrom((Description) field.getAnnotation(Description.class));
    }

    private static String descriptionFrom(Class<?> type) {
        return descriptionFrom((Description) type.getAnnotation(Description.class));
    }

    private static String descriptionFrom(Description description) {
        return description == null ? null : String.join(" ", description.value());
    }

    private static Class<?> getActualType(Type type) {
        if (type instanceof ParameterizedType) {
            ParameterizedType parameterizedType = (ParameterizedType) type;
            Type[] actualTypeArguments = parameterizedType.getActualTypeArguments();
            if (actualTypeArguments.length == 1) {
                return (Class<?>) actualTypeArguments[0];
            }
        }

        return null;
    }

    static boolean isCustomClass(Class<?> clazz) {
        if (clazz.getPackage() != null) {
            String packageName = clazz.getPackage().getName();
            if (packageName.startsWith("java.") || packageName.startsWith("javax.") || packageName.startsWith("jdk.") || packageName.startsWith("sun.") || packageName.startsWith("com.sun.")) {
                return false;
            }
        }

        return true;
    }

    public static Map<String, Map<String, Object>> toMap(Map<String, JsonSchemaElement> properties) {
        Map<String, Map<String, Object>> map = new LinkedHashMap<>();
        properties.forEach((property, value) -> {
            map.put(property, toMap(value));
        });
        return map;
    }

    public static Map<String, Object> toMap(JsonSchemaElement jsonSchemaElement) {
        LinkedHashMap<String, Object> properties;
        if (jsonSchemaElement instanceof JsonObjectSchema) {
            JsonObjectSchema jsonObjectSchema = (JsonObjectSchema) jsonSchemaElement;
            properties = new LinkedHashMap<>();
            properties.put("type", "object");
            if (jsonObjectSchema.description() != null) {
                properties.put("description", jsonObjectSchema.description());
            }

            properties.put("properties", toMap(jsonObjectSchema.properties()));
            if (jsonObjectSchema.required() != null) {
                properties.put("required", jsonObjectSchema.required());
            }

            return properties;
        } else if (jsonSchemaElement instanceof JsonArraySchema) {
            JsonArraySchema jsonArraySchema = (JsonArraySchema) jsonSchemaElement;
            properties = new LinkedHashMap<>();
            properties.put("type", "array");
            if (jsonArraySchema.description() != null) {
                properties.put("description", jsonArraySchema.description());
            }

            properties.put("items", toMap(jsonArraySchema.items()));
            return properties;
        } else if (jsonSchemaElement instanceof JsonEnumSchema) {
            JsonEnumSchema jsonEnumSchema = (JsonEnumSchema) jsonSchemaElement;
            properties = new LinkedHashMap<>();
            properties.put("type", "string");
            if (jsonEnumSchema.description() != null) {
                properties.put("description", jsonEnumSchema.description());
            }

            properties.put("enum", jsonEnumSchema.enumValues());
            return properties;
        } else if (jsonSchemaElement instanceof JsonStringSchema) {
            JsonStringSchema jsonStringSchema = (JsonStringSchema) jsonSchemaElement;
            properties = new LinkedHashMap<>();
            properties.put("type", "string");
            if (jsonStringSchema.description() != null) {
                properties.put("description", jsonStringSchema.description());
            }

            return properties;
        } else if (jsonSchemaElement instanceof JsonIntegerSchema) {
                JsonIntegerSchema jsonIntegerSchema = (JsonIntegerSchema) jsonSchemaElement;
            properties = new LinkedHashMap<>();
            properties.put("type", "integer");
            if (jsonIntegerSchema.description() != null) {
                properties.put("description", jsonIntegerSchema.description());
            }

            return properties;
        } else if (jsonSchemaElement instanceof JsonNumberSchema) {
            JsonNumberSchema jsonNumberSchema = (JsonNumberSchema) jsonSchemaElement;
            properties = new LinkedHashMap<>();
            properties.put("type", "number");
            if (jsonNumberSchema.description() != null) {
                properties.put("description", jsonNumberSchema.description());
            }

            return properties;
        } else if (jsonSchemaElement instanceof JsonBooleanSchema ) {
            JsonBooleanSchema jsonBooleanSchema = (JsonBooleanSchema) jsonSchemaElement;
            properties = new LinkedHashMap<>();
            properties.put("type", "boolean");
            if (jsonBooleanSchema.description() != null) {
                properties.put("description", jsonBooleanSchema.description());
            }

            return properties;
        } else {
            throw new IllegalArgumentException("Unknown type: " + String.valueOf(jsonSchemaElement.getClass()));
        }
    }

    public static class VisitedClassMetadata {
        public JsonSchemaElement jsonSchemaElement;
        public String reference;
        public boolean recursionDetected;

        public VisitedClassMetadata(JsonSchemaElement jsonSchemaElement, String reference, boolean recursionDetected) {
            this.jsonSchemaElement = jsonSchemaElement;
            this.reference = reference;
            this.recursionDetected = recursionDetected;
        }
    }
}
